import base64
import os
from io import BytesIO

import numpy as np
from PIL import Image, ImageOps
from loguru import logger
from ocr_process.rec.text_pipline import TextPipeline
from ocr_process.utils.img_utils import np2base64


class ImgParser(object):

    def __init__(self, det_model_path: str, rec_model_path: str, *, use_cuda=False, use_openvino=False):
        self.ocr_infer_model = TextPipeline(det_model_path=det_model_path, rec_model_path=rec_model_path, use_cuda=use_cuda, use_openvino=use_openvino)

    def parse(self, img):
        """
        parse
        """
        image = self.read_image(img)
        offset_x, offset_y = 0, 0
        img_h, img_w = image.shape[:2]
        res = {"image": np2base64(image), "offset_x": offset_x, "offset_y": offset_y, "img_w": img_w, "img_h": img_h}
        ocr_input = {"original_image": image, "image": image.copy()}
        ocr_result = self.ocr(ocr_input)
        logger.info(f"{ocr_result=}")
        # list_text = [item[1] for item in ocr_result]
        # logger.info(f"{','.join(list_text)}")
        res["layout"] = ocr_result
        return res

    def __call__(self, *args, **kwargs):
        """
        Call parse
        """
        return self.parse(*args, **kwargs)

    def ocr(self, image):
        """
        Call ocr for an image
        """

        def _get_box(box):
            box = [
                min(box[0][0], box[3][0]),  # x1
                min(box[0][1], box[1][1]),  # y1
                max(box[1][0], box[2][0]),  # x2
                max(box[2][1], box[3][1]),  # y2
            ]
            return box

        def _normal_box(box):
            # Ensure the height and width of bbox are greater than zero
            if box[3] - box[1] < 0 or box[2] - box[0] < 0:
                return False
            return True

        layout = []
        ocr_result = []
        dt_boxes, rec_res, _ = self.ocr_infer_model(image)
        if not dt_boxes and not rec_res:
            ocr_result.append(None)
        tmp_res = [[box.tolist(), res] for box, res in zip(dt_boxes, rec_res)]
        ocr_result.append(tmp_res)

        # ocr_result = self.ocr_infer_model.ocr(image)
        # print("ocr result:", ocr_result)
        ocr_result = ocr_result[0]
        if ocr_result:
            for segment in ocr_result:
                box = segment[0]
                box = _get_box(box)
                if not _normal_box(box):
                    continue
                text = segment[1][0]
                layout.append((box, text))
        return layout

    # @staticmethod
    # def _get_buffer(data):
    #     buff = None
    #     if len(data) < 1024:
    #         if os.path.exists(data):
    #             buff = open(data, "rb").read()
    #         else:
    #             raise FileNotFoundError("Image file {} not found!".format(data))
    #     if buff is None:
    #         buff = base64.b64decode(data)
    #     if buff:
    #         return BytesIO(buff)
    #     return buff

    @staticmethod
    def _get_buffer(data):
        # 使用os.path.abspath确保路径安全
        if len(data) < 1024 and os.path.exists(os.path.abspath(data)):
            try:
                with open(data, "rb") as file:
                    buff = file.read()
            except Exception as e:
                raise IOError(f"Error reading file {data}: {str(e)}")
        else:
            try:
                buff = base64.b64decode(data)
            except Exception as e:
                raise ValueError(f"Invalid base64 data: {str(e)}")
        return buff

    @classmethod
    def read_image(self, image):
        """
        read image to np.ndarray
        """
        if isinstance(image, np.ndarray):
            return image
        image_buff = self._get_buffer(image)

        # Use exif_transpose to correct orientation
        _image = np.array(ImageOps.exif_transpose(Image.open(BytesIO(image_buff)).convert("RGB")))
        return _image

    @staticmethod
    def _normalize_box(box, old_size, new_size, offset_x=0, offset_y=0):
        """normalize box"""
        return [
            int((box[0] + offset_x) * new_size[0] / old_size[0]),
            int((box[1] + offset_y) * new_size[1] / old_size[1]),
            int((box[2] + offset_x) * new_size[0] / old_size[0]),
            int((box[3] + offset_y) * new_size[1] / old_size[1]),
        ]
